# -*- coding: utf-8 -*-
"""
@project: picEnhance
@author: Wu Yue
@file: picEnhance
@ide: PyCharm
@creat time: 2020-09-01 08:36
@change time:
@function: 图片增强，将所有云图先处理为灰度图，再对 512*512 的图片进行竖直和横向翻转，旋转，保存
"""

import os, re, cv2
import numpy as np
import matplotlib.pyplot as plt
import shutil


""" 参数设置 """
# >>> path >>>
rootDir = "./"
sourceDirs = ["/grayPics"]
picSubDir = "/data"
labelSubDir = "/Label"
outputDir = "./outputPic"
# <<< path <<<

# >>> reguExp >>>
picRegu = "wind[0-9]*_[0-9]*.jpg"
labelRegu = "wind[0-9]*_[0-9]*.png"
# <<< reguExp <<<

# >>> file format >>>
picFormat = "jpg"
labelFormat = "png"
# <<< file format <<<
""" 参数设置 """


class singleEnhance:
    """
    对单个云图及其label进行增强 必须是 512*512 像素的图片
    """
    def __init__(self, picPath:str, labelPath:str, outputPicDir:str, outputLabelDir:str):
        """

        :param picPath: 单个云图完整的路径
        :param labelPath: 对应label完整的路径
        :param outputPicDir: 输出云图到的文件夹
        :param outputLabelDir: 输出label到的文件夹
        """
        self.picPath = picPath
        self.labelPath = labelPath
        self.outputPicDir = outputPicDir
        self.outputLabelDir = outputLabelDir

        self.pic = None  # 原始云图
        self.label = None  # 云图对应的 label

        # 处理操作
        print("正在处理: {} and {}".format(self.picPath, self.labelPath))
        self.changeToGray()
        self.enhance()

    def changeToGray(self):
        """
        如果原图是多通道的，则需要变为灰度图，替换原始文件
        :return:
        """
        print("\t灰度化处理原始图片...")

        self.pic = np.array(plt.imread(self.picPath))
        shutil.copyfile(self.labelPath, "./tempLabel.jpg")
        self.label = np.array(plt.imread("./tempLabel.jpg"))

        if len(self.pic.shape) == 3:  # RGBA
            r = self.pic[:, :, 0]
            g = self.pic[:, :, 1]
            b = self.pic[:, :, 2]
            a = self.pic[:, :, 3]

            assert np.equal(r, g).all() & np.equal(r, b).all()
            self.pic = r
            cv2.imwrite(self.picPath, self.pic)

        if len(self.label.shape) == 3:  # RGBA
            r = self.label[:, :, 0]
            g = self.label[:, :, 1]
            b = self.label[:, :, 2]
            a = self.label[:, :, 3]

            assert np.equal(r, g).all() & np.equal(r, b).all()
            self.label = r
            cv2.imwrite(self.labelPath, self.label)

        # 断言 图片尺寸都是512 * 512
        assert (self.pic.shape == (512, 512)) & (self.label.shape == (512, 512))

    def enhance(self):
        """
        图像增强
        :return:
        """
        print("\t图像扩充增强...")

        """原始图片"""
        newPicName = self.picPath.split("." + picFormat)[0].split("/")[-1] + "_0." + picFormat
        newLabelName = self.labelPath.split("." + labelFormat)[0].split("/")[-1] + "_0." + labelFormat
        newPic = self.pic
        newLabel = self.label
        cv2.imwrite(self.outputPicDir + "/" + newPicName, newPic)
        cv2.imwrite(self.outputLabelDir + "/" + newLabelName, newLabel)

        """上下翻转"""
        newPicName = self.picPath.split("." + picFormat)[0].split("/")[-1] + "_1." + picFormat
        newLabelName = self.labelPath.split("." + labelFormat)[0].split("/")[-1] + "_1." + labelFormat
        newPic = np.flip(self.pic, 0)
        newLabel = np.flip(self.label, 0)
        cv2.imwrite(self.outputPicDir + "/" + newPicName, newPic)
        cv2.imwrite(self.outputLabelDir + "/" + newLabelName, newLabel)

        """左右翻转"""
        newPicName = self.picPath.split("." + picFormat)[0].split("/")[-1] + "_2." + picFormat
        newLabelName = self.labelPath.split("." + labelFormat)[0].split("/")[-1] + "_2." + labelFormat
        newPic = np.flip(self.pic, 1)
        newLabel = np.flip(self.label, 1)
        cv2.imwrite(self.outputPicDir + "/" + newPicName, newPic)
        cv2.imwrite(self.outputLabelDir + "/" + newLabelName, newLabel)

        """逆时针旋转90°"""
        newPicName = self.picPath.split("." + picFormat)[0].split("/")[-1] + "_3." + picFormat
        newLabelName = self.labelPath.split("." + labelFormat)[0].split("/")[-1] + "_3." + labelFormat
        newPic = np.rot90(self.pic, k=1)  # k 代表旋转次数
        newLabel = np.rot90(self.label, k=1)
        cv2.imwrite(self.outputPicDir + "/" + newPicName, newPic)
        cv2.imwrite(self.outputLabelDir + "/" + newLabelName, newLabel)

        """逆时针旋转270°"""
        newPicName = self.picPath.split("." + picFormat)[0].split("/")[-1] + "_4." + picFormat
        newLabelName = self.labelPath.split("." + labelFormat)[0].split("/")[-1] + "_4." + labelFormat
        newPic = np.rot90(self.pic, k=3)
        newLabel = np.rot90(self.label, k=3)
        cv2.imwrite(self.outputPicDir + "/" + newPicName, newPic)
        cv2.imwrite(self.outputLabelDir + "/" + newLabelName, newLabel)


if __name__ == '__main__':
    for sourceDir in sourceDirs:
        for root, dirs, files in os.walk(rootDir + sourceDir + picSubDir):
            for file in files:
                if re.match(picRegu, file):
                    picPath = root + "/" + file
                    labelPath = rootDir + sourceDir + labelSubDir + "/" + file.split(".")[0] + "." + labelFormat
                    if os.path.exists(labelPath):
                        outputPicDir = outputDir + sourceDir + picSubDir
                        outputLabelDir = outputDir + sourceDir + labelSubDir
                        if not os.path.exists(outputPicDir): os.makedirs(outputPicDir)
                        if not os.path.exists(outputLabelDir): os.makedirs(outputLabelDir)
                        singleEnhance(picPath=picPath, labelPath=labelPath,
                                      outputPicDir=outputPicDir, outputLabelDir=outputLabelDir)


